import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 获取当前脚本所在目录
script_dir = os.path.dirname(os.path.abspath(__file__))
# 图像路径 - 使用绝对路径
image_path = os.path.join(script_dir, 'cat.jpg')

# 检查文件是否存在
if not os.path.exists(image_path):
    raise FileNotFoundError(f"图像文件不存在: {image_path}")

# 加载图像
img = tf.io.read_file(image_path)
img = tf.image.decode_image(img, channels=3)
img = tf.cast(img, tf.float32)
# 确保图像值在[0, 255]范围内
img = tf.clip_by_value(img, 0.0, 255.0)
img_array = img.numpy()

# 定义图像增强函数

def random_flip(image):
    """随机水平或垂直翻转"""
    # 50%的概率水平翻转
    image = tf.image.random_flip_left_right(image)
    # 50%的概率垂直翻转
    image = tf.image.random_flip_up_down(image)
    return image

def random_brightness(image, max_delta=0.2):
    """随机调整亮度"""
    image = tf.image.random_brightness(image, max_delta)
    return image

def random_contrast(image, lower=0.8, upper=1.2):
    """随机调整对比度"""
    image = tf.image.random_contrast(image, lower, upper)
    return image

def random_saturation(image, lower=0.5, upper=1.5):
    """随机调整饱和度"""
    image = tf.image.random_saturation(image, lower, upper)
    return image

# 创建一个图像增强管道
def augment_image(image):
    image = random_flip(image)
    image = random_brightness(image)
    image = random_contrast(image)
    image = random_saturation(image)
    # 确保输出值在[0,1]范围内
    image = tf.clip_by_value(image, 0.0, 1.0)
    return image

# 应用增强并显示结果

# 为了简化处理，我们使用tf.keras.preprocessing.image.ImageDataGenerator
def create_augmentation_generator():
    """创建图像增强生成器"""
    datagen = tf.keras.preprocessing.image.ImageDataGenerator(
        rotation_range=30,           # 随机旋转角度
        width_shift_range=0.2,       # 水平平移
        height_shift_range=0.2,      # 垂直平移
        shear_range=0.2,             # 剪切变换
        zoom_range=0.2,              # 随机缩放
        horizontal_flip=True,        # 水平翻转
        vertical_flip=True,          # 垂直翻转
        brightness_range=[0.5, 1.5], # 亮度调整
        fill_mode='nearest'          # 填充模式
    )
    return datagen

# 加载图像并调整大小
# 确保图像在处理前值在[0, 255]范围内
img = tf.clip_by_value(img, 0.0, 255.0)
img_resized = tf.image.resize(img, [224, 224])
# 归一化到[0,1]范围用于处理
img_resized = img_resized / 255.0
# 确保归一化后值在[0,1]范围内
img_resized = tf.clip_by_value(img_resized, 0.0, 1.0)
print(f"调整大小后图像形状: {img_resized.shape}, 数据类型: {img_resized.dtype}, 值范围: [{tf.reduce_min(img_resized)}, {tf.reduce_max(img_resized)}]")

# 扩展维度以适应ImageDataGenerator
img_batch = tf.expand_dims(img_resized, 0)

# 生成多个增强样本
datagen = create_augmentation_generator()
augmented_images = []

# 生成增强图像
for i, batch in enumerate(datagen.flow(img_batch, batch_size=1)):
    # 确保生成的图像值在[0,1]范围内
    augmented_img = tf.clip_by_value(batch[0], 0.0, 1.0)
    augmented_images.append(augmented_img)
    if i >= 5:  # 生成6个样本（包括原始图像）
        break

# 显示原始图像和增强后的图像
plt.figure(figsize=(15, 10))

# 显示原始图像
plt.subplot(2, 4, 1)
# 确保原始图像值在[0,1]范围内
img_display = tf.clip_by_value(img_resized, 0.0, 1.0)
plt.imshow(img_display)
plt.title("原始图像")
plt.axis('off')

# 显示增强后的图像
for i, aug_img in enumerate(augmented_images):
    plt.subplot(2, 4, i+2)
    # 确保增强图像值在[0,1]范围内
    aug_img_display = tf.clip_by_value(aug_img, 0.0, 1.0)
    plt.imshow(aug_img_display)
    plt.title(f"增强图像 {i+1}")
    plt.axis('off')

plt.tight_layout()
plt.savefig(os.path.join(script_dir, 'tf_augmentation_result.png'))
plt.show()

# 使用tf.image进行另一种增强方式演示
def demonstrate_tf_image_augmentation():
    """使用tf.image进行图像增强演示"""
    print("使用tf.image进行图像增强演示...")
    
    # 创建一个包含原始图像的批次
    batch_img = tf.expand_dims(img_resized, 0)
    
    # 应用各种增强
    flipped_lr = tf.image.flip_left_right(img_resized)
    flipped_ud = tf.image.flip_up_down(img_resized)
    rotated_90 = tf.image.rot90(img_resized)
    adjusted_brightness = tf.image.adjust_brightness(img_resized, 0.3)
    adjusted_contrast = tf.image.adjust_contrast(img_resized, 1.5)
    adjusted_saturation = tf.image.adjust_saturation(img_resized, 1.5)
    
    # 显示结果
    plt.figure(figsize=(15, 12))
    
    images = [
        img_resized, flipped_lr, flipped_ud, rotated_90,
        adjusted_brightness, adjusted_contrast, adjusted_saturation
    ]
    
    titles = [
        "原始图像", "水平翻转", "垂直翻转", "旋转90度",
        "亮度调整", "对比度调整", "饱和度调整"
    ]
    
    for i, (image, title) in enumerate(zip(images, titles)):
        plt.subplot(2, 4, i+1)
        # 确保图像值在[0,1]范围内
        image_display = tf.clip_by_value(image, 0.0, 1.0)
        plt.imshow(image_display)
        plt.title(title)
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(os.path.join(script_dir, 'tf_image_augmentation_result.png'))
    plt.show()

# 运行tf.image增强演示
demonstrate_tf_image_augmentation()

print("基于TensorFlow的图像增强演示完成！")
print("结果已保存为 'tf_augmentation_result.png' 和 'tf_image_augmentation_result.png'")